// ==UserScript==
// @name         bliveproxy
// @namespace    https://github.com/xfgryujk/bliveproxy
// @version      0.5.0.1
// @author       xfgryujk
// @description  B站直播websocket hook框架
// @match        *://*/*
// @grant        none
// @run-at       document-start
// ==/UserScript==

// 使用方法：
// bliveproxy.addCommandHandler('DANMU_MSG', command => {
//   console.log(command)
//   let info = command.info
//   info[1] = '测试'
// })
//
// 如果@grant不是none，则要使用unsafeWindow.bliveproxy

(function (W) {
  const HEADER_SIZE = 16

  const WS_BODY_PROTOCOL_VERSION_NORMAL = 0
  const WS_BODY_PROTOCOL_VERSION_HEARTBEAT = 1
  const WS_BODY_PROTOCOL_VERSION_BROTLI = 3

  const OP_HEARTBEAT_REPLY = 3 // WS_OP_HEARTBEAT_REPLY
  const OP_SEND_MSG_REPLY = 5 // WS_OP_MESSAGE
  const OP_AUTH_REPLY = 8 // WS_OP_CONNECT_SUCCESS

  // 兼容@grant不是none的情况
  let pageWindow = W
  // let BrotliDecode = (await import('https://cdn.jsdelivr.net/gh/google/brotli@f8c671774514357abec2f6b14c8ee13c6fd885d3/js/decode.min.js')).BrotliDecode

  let textEncoder = new TextEncoder()
  let textDecoder = new TextDecoder()

  function main() {
    if (pageWindow.bliveproxy) {
      // 防止多次加载
      return
    }
    initApi()
    hook()
  }

  function initApi() {
    pageWindow.bliveproxy = api
  }

  let api = {
    addCommandHandler(cmd, handler) {
      let handlers = this._commandHandlers[cmd]
      if (!handlers) {
        handlers = this._commandHandlers[cmd] = []
      }
      handlers.push(handler)
    },
    removeCommandHandler(cmd, handler) {
      let handlers = this._commandHandlers[cmd]
      if (!handlers) {
        return
      }
      this._commandHandlers[cmd] = handlers.filter(item => item !== handler)
    },

    // 私有API
    _commandHandlers: {},
    _getCommandHandlers(cmd) {
      return this._commandHandlers[cmd] || null
    }
  }

  function hook() {
    pageWindow.WebSocket = new Proxy(pageWindow.WebSocket, {
      construct(target, args) {
        let obj = new target(...args)
        return new Proxy(obj, proxyHandler)
      }
    })
  }

  let proxyHandler = {
    get(target, property) {
      let value = target[property]
      if ((typeof value) === 'function') {
        value = value.bind(target)
      }
      return value
    },
    set(target, property, value) {
      if (property === 'onmessage') {
        let realOnMessage = value
        value = function (event) {
          myOnMessage(event, realOnMessage)
        }
      }
      target[property] = value
      return value
    }
  }

  function myOnMessage(event, realOnMessage) {
    if (!(event.data instanceof ArrayBuffer)) {
      realOnMessage(event)
      return
    }

    let data = new Uint8Array(event.data)
    function callRealOnMessageByPacket(packet) {
      realOnMessage({ ...event, data: packet })
    }
    handleMessage(data, callRealOnMessageByPacket)
  }

  function makePacketFromCommand(command) {
    let body = textEncoder.encode(JSON.stringify(command))
    return makePacketFromUint8Array(body, OP_SEND_MSG_REPLY)
  }

  function makePacketFromUint8Array(body, operation) {
    let packLen = HEADER_SIZE + body.byteLength
    let packet = new ArrayBuffer(packLen)

    // 不需要压缩
    let ver = operation === OP_HEARTBEAT_REPLY ? WS_BODY_PROTOCOL_VERSION_HEARTBEAT : WS_BODY_PROTOCOL_VERSION_NORMAL
    let packetView = new DataView(packet)
    packetView.setUint32(0, packLen)        // pack_len
    packetView.setUint16(4, HEADER_SIZE)    // raw_header_size
    packetView.setUint16(6, ver)            // ver
    packetView.setUint32(8, operation)      // operation
    packetView.setUint32(12, 1)             // seq_id

    let packetBody = new Uint8Array(packet, HEADER_SIZE, body.byteLength)
    for (let i = 0; i < body.byteLength; i++) {
      packetBody[i] = body[i]
    }
    return packet
  }

  function handleMessage(data, callRealOnMessageByPacket) {
    let dataView = new DataView(data.buffer)
    let operation = dataView.getUint32(8)

    switch (operation) {
      case OP_AUTH_REPLY:
      case OP_SEND_MSG_REPLY: {
        // 可能有多个包一起发，需要分包
        let offset = 0
        while (offset < data.byteLength) {
          let dataView = new DataView(data.buffer, offset)
          let packLen = dataView.getUint32(0)
          let rawHeaderSize = dataView.getUint16(4)
          let ver = dataView.getUint16(6)
          let operation = dataView.getUint32(8)
          // let seqId = dataView.getUint32(12)

          let body = new Uint8Array(data.buffer, offset + rawHeaderSize, packLen - rawHeaderSize)
          if (operation === OP_SEND_MSG_REPLY) {
            // 业务消息
            switch (ver) {
              case WS_BODY_PROTOCOL_VERSION_NORMAL: {
                // body是单个JSON消息
                body = textDecoder.decode(body)
                body = JSON.parse(body)
                handleCommand(body, callRealOnMessageByPacket)
                break
              }
              case WS_BODY_PROTOCOL_VERSION_BROTLI: {
                // body是压缩过的多个消息
                body = pageWindow.BrotliDecode(body)
                handleMessage(body, callRealOnMessageByPacket)
                break
              }
              default: {
                // 未知的body格式
                let packet = makePacketFromUint8Array(body, operation)
                callRealOnMessageByPacket(packet)
                break
              }
            }
          } else {
            // 非业务消息
            let packet = makePacketFromUint8Array(body, operation)
            callRealOnMessageByPacket(packet)
          }

          offset += packLen
        }
        break
      }

      // 服务器心跳包，前4字节是人气值，后面是客户端发的心跳包内容
      // packLen不包括客户端发的心跳包内容，不知道是不是服务器BUG
      // 这里没用到心跳包就不处理了
      // case OP_HEARTBEAT_REPLY:
      default: {
        // 只有一个包
        let packLen = dataView.getUint32(0)
        let rawHeaderSize = dataView.getUint16(4)

        let body = new Uint8Array(data.buffer, rawHeaderSize, packLen - rawHeaderSize)
        let packet = makePacketFromUint8Array(body, operation)
        callRealOnMessageByPacket(packet)
        break
      }
    }
  }

  function handleCommand(command, callRealOnMessageByPacket) {
    if (command instanceof Array) {
      for (let oneCommand of command) {
        this.handleCommand(oneCommand)
      }
      return
    }

    let cmd = command.cmd || ''
    let pos = cmd.indexOf(':')
    if (pos != -1) {
      cmd = cmd.substr(0, pos)
    }
    let handlers = api._getCommandHandlers(cmd)
    if (handlers) {
      for (let handler of handlers) {
        handler(command)
      }
    }
    // 全部监控
    let ALLhandlers = api._getCommandHandlers('ALL')
    if (ALLhandlers) {
      for (let ALLhandler of ALLhandlers) {
        ALLhandler(command)
      }
    }
    // console.log(command)

    let packet = makePacketFromCommand(command)
    callRealOnMessageByPacket(packet)
  }

  main()
})(typeof unsafeWindow === 'undefined' ? window : unsafeWindow);